package cc.gpai.data_stru.bag;

import java.util.AbstractSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;

public abstract class MapBag<E> extends AbstractBag<E> {

	Map<E, AtomicLong> map = createMap();

	protected abstract Map<E, AtomicLong> createMap();

	@Override
	public boolean isEmpty() {
		return map.isEmpty();
	}

	@Override
	public boolean contains(Object o) {
		return map.containsKey(o);
	}

	@Override
	public boolean add(E e) {
		AtomicLong l = map.get(e);
		if (l == null) {
			map.put(e, new AtomicLong(1));
		} else {
			l.incrementAndGet();
		}
		return true;
	}

	@Override
	public long modCount(E e, long count) {
		AtomicLong l = map.get(e);
		if (l != null) {
			long after = l.addAndGet(count);
			if (after <= 0) {
				map.remove(e);
				return 0;
			}
			return after;
		} else {
			if (count > 0) {
				map.put(e, new AtomicLong(count));
				return count;
			} else {
				return 0;
			}
		}
	}

	@Override
	public void setCount(E e, long count) {
		if (count < 1) {
			remove(e);
			return;
		}
		AtomicLong l = map.get(e);
		if (l == null) {
			map.put(e, new AtomicLong(count));
		} else {
			l.set(count);
		}
	}

	@Override
	public boolean remove(Object o) {
		long count = getCount(o);
		if (count == 0) {
			return false;
		} else if (count == 1) {
			map.remove(o);
		} else {
			map.get(o).decrementAndGet();
		}
		return true;
	}

	@Override
	public void clear() {
		map.clear();
	}

	@Override
	public long getCount(Object obj) {
		AtomicLong l = map.get(obj);
		return l == null ? 0 : l.longValue();
	}

	@Override
	public Set<E> uniqueSet() {
		return new AbstractSet<E>() {

			@Override
			public boolean add(E e) {
				return MapBag.this.add(e);
			}

			@Override
			public int size() {
				return map.size();
			}

			@Override
			public Iterator<E> iterator() {
				return map.keySet().iterator();
			}
		};
	}

	@Override
	public int size() {
		int sum = 0;
		for (AtomicLong i : map.values()) {
			sum += i.get();
		}
		return sum;
	}

	@Override
    public Set<Entry<E, AtomicLong>> entrySet() {
		return map.entrySet();
	}

}